#!/usr/bin/env python
# encoding: utf-8
"""
@author: zhifeng
@license: (C) Copyright 2013-2017, Vivo Corporation Limited.
@contact: zhifeng.ding@vivo.com
@software: image vision
@file: mobilenetv3_segmentation.py
@time: 2020/7/16 16:37
@desc:
"""
import tensorflow as tf
from .MobileNetV3Small import MobileNetV3Small
from .literassp import LiteRASSP


class MobileNetV3Segment(tf.keras.Model):
    def __init__(self, num_classes, **kwargs):
        super(MobileNetV3SmallSeg, self).__init__(**kwargs)

        backbone = MobileNetV3Small(num_classes=num_classes)
        self.backbones = []

        self.names = ['feature_1_8', 'feature_1_16']
        for layer in backbone.layers:
            self.backbones.append(layer)
            if layer.name == self.names[-1]:
                break

        self.seg_head = LiteRASSP(num_classes=num_classes)

    def call(self, inputs, training=None, mask=None):
        x = inputs
        outputs = []
        for layer in self.backbones:
            x = layer(x, training=training)
            if layer.name in self.names:
                outputs.append(x)

        seg_out = self.seg_head(outputs)
        return seg_out


if __name__ == '__main__':
    model = MobileNetV3SmallSeg(num_classes=10)
    x = tf.random.normal((2, 224, 224, 3))
    y = model(x)
    print(y.shape)
